#!/usr/bin/env python3
"""
配置AI模型到MongoDB数据库
"""

import asyncio
import os
from motor.motor_asyncio import AsyncIOMotorClient
from datetime import datetime

# MongoDB连接配置
MONGODB_URL = os.getenv("MONGODB_URL", "mongodb://admin:newhand@localhost:27017/")
DATABASE_NAME = os.getenv("MONGODB_DATABASE", "ai_writing")

async def setup_ai_models():
    """配置AI模型"""

    # 连接MongoDB
    client = AsyncIOMotorClient(MONGODB_URL)
    db = client[DATABASE_NAME]
    collection = db.model_configs

    print("=" * 60)
    print("AI模型配置工具")
    print("=" * 60)

    # 清理旧配置
    deleted = await collection.delete_many({})
    print(f"清理了 {deleted.deleted_count} 个旧配置")

    # 配置模型列表
    models = [
        {
            "id": "deepseek-chat",
            "name": "DeepSeek Chat",
            "provider": "deepseek",
            "enabled": True,
            "status": "healthy",
            "api_key": "sk-4f5f87fa95dd490e8f19a014f37b3ac8",  # DeepSeek API Key
            "api_base": "https://api.deepseek.com",
            "model_name": "deepseek-chat",
            "description": "DeepSeek Chat 模型，高性价比的中文对话模型",
            "max_tokens": 4096,
            "temperature": 0.7,
            "created_at": datetime.utcnow(),
            "updated_at": datetime.utcnow()
        },
        {
            "id": "qwen-turbo",
            "name": "通义千问 Turbo",
            "provider": "qwen",
            "enabled": True,
            "status": "healthy",
            "api_key": "sk-your-qwen-api-key",  # 需要替换为实际的API Key
            "api_base": "https://dashscope.aliyuncs.com/compatible-mode/v1",
            "model_name": "qwen-turbo",
            "description": "阿里通义千问模型",
            "max_tokens": 4096,
            "temperature": 0.7,
            "created_at": datetime.utcnow(),
            "updated_at": datetime.utcnow()
        },
        {
            "id": "gpt-3.5-turbo",
            "name": "GPT-3.5 Turbo",
            "provider": "openai",
            "enabled": False,  # 默认禁用，需要有效的API Key
            "status": "disabled",
            "api_key": "sk-your-openai-api-key",  # 需要替换为实际的API Key
            "api_base": "https://api.openai.com",
            "model_name": "gpt-3.5-turbo",
            "description": "OpenAI GPT-3.5 模型",
            "max_tokens": 4096,
            "temperature": 0.7,
            "created_at": datetime.utcnow(),
            "updated_at": datetime.utcnow()
        }
    ]

    # 插入模型配置
    result = await collection.insert_many(models)
    print(f"\n✅ 成功插入 {len(result.inserted_ids)} 个模型配置")

    # 显示配置的模型
    print("\n已配置的模型：")
    print("-" * 60)
    async for model in collection.find({"enabled": True}):
        print(f"  • {model['name']} ({model['id']})")
        print(f"    Provider: {model['provider']}")
        print(f"    API Base: {model['api_base']}")
        print(f"    Status: {model['status']}")
        print()

    # 测试提示生成
    print("测试提示生成...")
    print("-" * 60)

    # 获取一个启用的模型
    test_model = await collection.find_one({"enabled": True, "status": "healthy"})

    if test_model:
        print(f"使用模型: {test_model['name']}")
        print(f"API Key: ***{test_model['api_key'][-4:]}")

        # 测试API调用
        import aiohttp

        # 构建端点
        api_base = test_model['api_base'].rstrip('/')
        if test_model['provider'] == 'deepseek':
            endpoint = f"{api_base}/v1/chat/completions"
        else:
            endpoint = f"{api_base}/v1/chat/completions"

        payload = {
            "model": test_model['model_name'],
            "messages": [
                {"role": "system", "content": "你是一个测试助手"},
                {"role": "user", "content": "回复OK"}
            ],
            "max_tokens": 10,
            "temperature": 0.1
        }

        headers = {
            "Authorization": f"Bearer {test_model['api_key']}",
            "Content-Type": "application/json"
        }

        print(f"测试端点: {endpoint}")

        try:
            async with aiohttp.ClientSession() as session:
                async with session.post(
                    endpoint,
                    headers=headers,
                    json=payload,
                    timeout=aiohttp.ClientTimeout(total=5)
                ) as response:
                    if response.status == 200:
                        result = await response.json()
                        print("✅ API测试成功!")
                        if 'choices' in result:
                            content = result['choices'][0]['message']['content']
                            print(f"响应: {content}")
                    else:
                        error = await response.text()
                        print(f"❌ API错误 {response.status}: {error[:100]}")
        except aiohttp.ClientTimeout:
            print("⏱️ API调用超时")
        except Exception as e:
            print(f"❌ 连接错误: {e}")

    print("\n" + "=" * 60)
    print("配置完成!")
    print("=" * 60)
    print("\n提示：")
    print("1. DeepSeek模型已配置并启用")
    print("2. 如需使用其他模型，请替换相应的API Key")
    print("3. 重启AI服务以应用新配置")

    # 关闭连接
    client.close()

if __name__ == "__main__":
    asyncio.run(setup_ai_models())